from pathlib import Path
import numpy as np
import cv2

__all__ = ['cv_show', 'SkullTemplateMatch', 'draw_contours_to_mask', 'draw_vetebra_landmarks', 'cv_keyboard_run',
           'cv_resize_shape', 'convert_mask_to_image', 'cv_imread', 'cv_cvtcolor', 'cv_imwrite']

vetebra_colors = [[0.76590096, 0.0266074, 0.9806378],
           [0.54197179, 0.81682527, 0.95081629],
           [0.0799733, 0.79737015, 0.15173816],
           [0.93240442, 0.8993321, 0.09901344],
           [0.73130136, 0.05366301, 0.98405681],
           [0.01664966, 0.16387004, 0.94158259],
           [0.54197179, 0.81682527, 0.45081629],
           # [0.92074915, 0.09919099 ,0.97590748],
           [0.83445145, 0.97921679, 0.12250426],
           [0.7300924, 0.23253621, 0.29764521],
           [0.3856775, 0.94859286, 0.9910683],  # 10
           [0.45762137, 0.03766411, 0.98755338],
           [0.99496697, 0.09113071, 0.83322314],
           [0.96478873, 0.0233309, 0.13149931],
           [0.33240442, 0.9993321 , 0.59901344],
            # [0.77690519,0.81783954,0.56220024],
           # [0.93240442, 0.8993321, 0.09901344],
           [0.95815068, 0.88436046, 0.55782268],
           [0.03728425, 0.0618827, 0.88641827],
           [0.05281129, 0.89572238, 0.08913828],

           ]


def draw_vetebra_landmarks(pts, image, const_color=None):
    for i in range(pts.shape[0]//4):
        # color = np.random.rand(3)
        pt = pts[i*4: i*4 + 4]
        if const_color is not None:
            color_255 = const_color
        else:
            color = vetebra_colors[i]
            color_255 = (255 * color[0], 255 * color[1], 255 * color[2])
        # cv2.circle(ori_image, (int(pt[0]), int(pt[1])), 6, (255,255,255), -1,1)
        cv2.circle(image, (int(pt[0][0]), int(pt[0][1])), 5, color_255, -1, 1)
        cv2.circle(image, (int(pt[1][0]), int(pt[1][1])), 5, color_255, -1, 1)
        cv2.circle(image, (int(pt[2][0]), int(pt[2][1])), 5, color_255, -1, 1)
        cv2.circle(image, (int(pt[3][0]), int(pt[3][1])), 5, color_255, -1, 1)
    return image


def cv_show(name, image, resize=1):
    H,W = image.shape[0:2]
    cv2.namedWindow(name, cv2.WINDOW_NORMAL)
    cv2.imshow(name, image.astype(np.uint8))
    cv2.resizeWindow(name, round(resize*W), round(resize*H))


def cv_resize_shape(image, shape, interpolation=cv2.INTER_LINEAR):
    """
    INTER_AREA = 3
    INTER_BITS = 5
    INTER_BITS2 = 10
    INTER_CUBIC = 2
    INTER_LANCZOS4 = 4
    INTER_LINEAR = 1
    :param image:
    :param shape:
    :param interpolation:
    :return:
    """
    assert len(shape) == 2
    return cv2.resize(image, shape[::-1], interpolation=interpolation)


def convert_mask_to_image(img:np.ndarray, mask:np.ndarray, classes=2, palette=None):
    if palette is None:
        palette = np.random.randint(0, 255, size=(classes, 3))
    else:
        palette = np.array(palette)
    color_seg = np.zeros((img.shape[0], img.shape[1], 3), dtype=np.uint8)

    if mask.shape[0] != img.shape[0] or mask.shape[1] != img.shape[1]:
        mask = cv2.resize(mask,  (img.shape[1],  img.shape[0]), cv2.INTER_LINEAR)

    for label, color in enumerate(palette):
        color_seg[mask == label, :] = color
    # convert to BGR
    color_seg = color_seg[..., ::-1]

    img = img * 0.5 + color_seg * 0.5
    img = img.astype(np.uint8)
    return img, color_seg


def cv_imread(filename, flags=0):
    """
    IMREAD_ANYCOLOR = 4
    IMREAD_ANYDEPTH = 2
    IMREAD_COLOR = 1
    IMREAD_GRAYSCALE = 0
    :param filename:
    :param flags:
    :return:
    """
    return cv2.imread(filename, flags=flags)

def cv_imwrite(filename, image):
    """
    IMREAD_ANYCOLOR = 4
    IMREAD_ANYDEPTH = 2
    IMREAD_COLOR = 1
    IMREAD_GRAYSCALE = 0
    :param filename:
    :param flags:
    :return:
    """
    return cv2.imwrite(filename, image)


def cv_cvtcolor(image, code=cv2.COLOR_BGR2GRAY):
    return cv2.cvtColor(image, code=code)

class SkullTemplateMatch(object):
    def __init__(self, template:np.ndarray, method=cv2.TM_CCOEFF, skull_ratio=0.35):
        """
        平方差匹配CV_TM_SQDIFF：用两者的平方差来匹配，最好的匹配值为0
        归一化平方差匹配CV_TM_SQDIFF_NORMED
        相关匹配CV_TM_CCORR：用两者的乘积匹配，数值越大表明匹配程度越好
        归一化相关匹配CV_TM_CCORR_NORMED
        相关系数匹配CV_TM_CCOEFF：用两者的相关系数匹配，1表示完美的匹配，-1表示最差的匹配
        归一化相关系数匹配CV_TM_CCOEFF_NORMED
        :param method:
        """
        self.method = method
        self.template = template
        self.skull_ratio = skull_ratio

    def match_image(self, image):
        res = cv2.matchTemplate(image, self.template, self.method)
        res = cv2.normalize(res, None, 0, 1, cv2.NORM_MINMAX)
        min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)

        h_template, w_template = self.template.shape[0], self.template.shape[1]
        left_top = max_loc  # 左上角
        right_bottom = (left_top[0] + w_template, left_top[1] + h_template-10)  # 右下角

        skull_yes = False
        skull_ratio_by_height = right_bottom[1] / image.shape[0]
        if skull_ratio_by_height < self.skull_ratio:
            skull_yes = True
        return (left_top, right_bottom), skull_yes


def clip(img, dtype, maxval):
    return np.clip(img, 0, maxval).astype(dtype)


def _equalize_cv(img, mask=None):
    """

    :param img: [h, w] gray image
    :param mask:
    :return:
    """
    if mask is None:
        return cv2.equalizeHist(img)

    histogram = cv2.calcHist([img], [0], mask, [256], (0, 256)).ravel()
    i = 0
    for val in histogram:
        if val > 0:
            break
        i += 1
    i = min(i, 255)

    total = np.sum(histogram)
    if histogram[i] == total:
        return np.full_like(img, i)

    scale = 255.0 / (total - histogram[i])
    _sum = 0

    lut = np.zeros(256, dtype=np.uint8)
    i += 1
    for i in range(i, len(histogram)):
        _sum += histogram[i]
        lut[i] = clip(round(_sum * scale), np.dtype("uint8"), 255)

    return cv2.LUT(img, lut)

def draw_contours_to_mask(mask, cnt:list):
    """
    cv2 drawContours and generate mask
    :param mask:
    :param cnt:
    :return:
    """
    draw_image = np.zeros_like(mask)
    draw_image = cv2.cvtColor(draw_image, cv2.COLOR_GRAY2BGR)
    cv2.drawContours(draw_image, cnt, 0, (255, 255, 255), -1)
    # cv_show('draw', draw_image, 0.5)
    # cv2.waitKey(0)
    draw_image = (cv2.cvtColor(draw_image, cv2.COLOR_BGR2GRAY) > 0).astype(np.uint8)

    return draw_image


def cv_keyboard_run(key='q'):
    k = cv2.waitKey(0) & 0xFF
    if k == ord(key):
        cv2.destroyAllWindows()

def main():


    root_dir = Path('/home/blake/data/medical/datasets/vertebral/spine_full_length_photography2')
    images = root_dir.rglob("*.jpg")

    template = cv2.imread('image/skull.bmp', cv2.IMREAD_GRAYSCALE)

    skull_matcher = SkullTemplateMatch(template, method=cv2.TM_CCOEFF)

    template = cv2.imread('image/Pelvis.bmp', cv2.IMREAD_GRAYSCALE)

    #pelvis_matcher = SkullTemplateMatch(template, method=cv2.TM_CCOEFF)

    for file in images:
        image = cv2.imread(str(file), cv2.IMREAD_GRAYSCALE)
        #image = cv2.equalizeHist(image)
        res1, skull_yes = skull_matcher.match_image(image)
        show_image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
        #res2 = pelvis_matcher.match_image(image)
        cv2.rectangle(show_image, res1[0], res1[1], (0, 0, 255), 2)
        if skull_yes:
            cv2.putText(show_image, "YES,Skull!", (50, 50), cv2.FONT_HERSHEY_COMPLEX, 2, (0, 255, 0), 2)
        else:
            cv2.putText(show_image, "Ohp,Nothing!", (50, 50), cv2.FONT_HERSHEY_COMPLEX, 2, (0, 255, 0), 2)

        cv_show(file.name, show_image, 1.0)
        cv2.waitKey(0)
        cv2.destroyWindow(file.name)



if __name__ == '__main__':
    main()
